import re
from transformers import AutoTokenizer, AutoModel
import torch
from pymilvus import connections, Collection, db
import gradio as gr

# 初始化模型和tokenizer
tokenizer = AutoTokenizer.from_pretrained("microsoft/BiomedNLP-BiomedBERT-large-uncased-abstract")
model = AutoModel.from_pretrained("microsoft/BiomedNLP-BiomedBERT-large-uncased-abstract")

# 配置
DATABASE_NAME = "my_medical_db"
COLLECTION_NAME = "new_medical_docs"
DIMENSION = 1024


# 连接Milvus
def connect_milvus():
    connections.connect(host='localhost', port='19530')
    db.using_database(DATABASE_NAME)


# 生成查询嵌入
def generate_query_embedding(query_text):
    inputs = tokenizer(query_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
        embedding = outputs.last_hidden_state[:, 0, :].squeeze().numpy().tolist()
    return embedding


# 高亮关键词（HTML格式）
def highlight_keywords(text, keywords):
    for keyword in keywords:
        text = re.sub(f"(?i)({keyword})", r"<b>\1</b>", text)  # 用<b>高亮
    return text


# 语义搜索
def semantic_search(collection, query_text, top_k):
    query_embedding = generate_query_embedding(query_text)
    collection.load()
    search_params = {"metric_type": "L2", "params": {"nprobe": 10}}
    results = collection.search(
        data=[query_embedding],
        anns_field="embedding",
        param=search_params,
        limit=top_k,
        output_fields=["id", "chunk_text"]
    )
    return [(result.entity.get("id"), result.entity.get("chunk_text"), 1 / (1 + result.distance)) for result in
            results[0]]


# 关键词搜索
def keyword_search(collection, keywords, top_k):
    collection.load()
    expr = " || ".join(f"chunk_text like '%{kw}%'" for kw in keywords)
    results = collection.query(expr=expr, limit=top_k, output_fields=["id", "chunk_text"])
    return [(res["id"], res["chunk_text"], 0.8) for res in results]


# 混合检索（OR关系）
def hybrid_search(collection, query_text, keywords, top_k=5):
    results_dict = {}

    if query_text:
        semantic_results = semantic_search(collection, query_text, top_k)
        for id_, text, score in semantic_results:
            results_dict[id_] = {"id": id_, "text": text, "score": score}

    if keywords:
        keyword_results = keyword_search(collection, keywords, top_k)
        for id_, text, score in keyword_results:
            if id_ in results_dict:
                results_dict[id_]["score"] = max(results_dict[id_]["score"], score)
            else:
                results_dict[id_] = {"id": id_, "text": text, "score": score}

    final_results = []
    for item in results_dict.values():
        highlighted_text = highlight_keywords(item["text"], keywords) if keywords else item["text"]
        final_results.append({"id": item["id"], "text": highlighted_text, "score": item["score"]})

    return sorted(final_results, key=lambda x: x["score"], reverse=True)[:top_k]


# Gradio搜索函数
def search(query_text, keywords_input):
    if not query_text and not keywords_input:
        return "请输入查询内容或关键词"

    connect_milvus()
    collection = Collection(COLLECTION_NAME)

    keywords = keywords_input.split() if keywords_input else []
    results = hybrid_search(collection, query_text, keywords)

    if not results:
        return "未找到匹配结果"

    output = f"查询: {query_text or '无'}, 关键词: {keywords or '无'}<br><br>"
    for result in results:
        output += f"ID: {result['id']}<br>"
        output += f"Score: {result['score']:.4f}<br>"
        output += f"Text: {result['text']}<br>"
        output += "-" * 50 + "<br>"
    return output


# Gradio界面
with gr.Blocks(title="医学文献搜索") as demo:
    gr.Markdown("# 医学文献搜索系统")
    with gr.Row():
        query_input = gr.Textbox(label="查询内容（可选）", placeholder="请输入查询内容...")
        keywords_input = gr.Textbox(label="关键词（可选，用空格分隔）", placeholder="请输入关键词...")
    search_button = gr.Button("搜索")
    output = gr.HTML(label="搜索结果")

    search_button.click(fn=search, inputs=[query_input, keywords_input], outputs=output)

demo.launch()